from numpy import *
from io import *
import copy
import random
import os
import pylab 

# poceedure that makes a n point running average of the given list or array
def running_average(l,n):
    lseg = []
    lreturn = []
    for i in range(len(l)):
        if i < n:
            lseg.append(l[i])
            lreturn.append(sum(lseg)/len(lseg))
        else:
            lseg = lseg[1:]
            lseg.append(l[i])
            lreturn.append(sum(lseg)/len(lseg))
    return lreturn

def resid(data,fit):
    rreturn = 0
    for i in range(len(data)):
        rreturn += (fit[i] - data[i])**2
    return rreturn

class synapse_info:
    def __init__(self):
        dict = read_dict_float("configure")
        self.__dict__.update(dict)
        # the range over which the parameter may vary ( suffix W )
        self.EonsetW   = self.Eonsets[-1]  - self.Eonsets[0]
        self.EtauW     = self.Etaus[-1]    - self.Etaus[0]
        self.EgmaxW    = self.Egmaxs[-1]   - self.Egmaxs[0]
        self.ErevW     = self.Erevs[-1]    - self.Erevs[0]
        self.IonsetW   = self.Ionsets[-1]  - self.Ionsets[0]
        self.ItauW     = self.Itaus[-1]    - self.Itaus[0]
        self.IgmaxW    = self.Igmaxs[-1]   - self.Igmaxs[0]
        self.IrevW     = self.Irevs[-1]    - self.Irevs[0]
        self.gpasW     = self.gpass[-1]    - self.gpass[0]

        infile = open("OS500MSHYPE.DAT","r")
        lines = infile.readlines()
        lines = lines[5020:10020]
        # read the data in
        self.data = []
        self.data.append( traces([0,2,4],lines) )
        self.data.append( traces([6,8,10],lines) )
        self.data.append( traces([12,14,16],lines) )
        self.data.append( traces([18,20,22],lines) )
        self.data.append( traces([24,26,28],lines) )
        self.holding_potentials = [-54.9,-67.4,-71.6,-83,-90.6]
        self.time = arange(0,500,0.1)
        self.image_counter = 0

    def neuron_print(self,holding_potential):
        ofile = open("neuron.out","w")
        ofile.write("0.1\t500\n")
        ofile.write("%f\t%f\t%f\t%f\n" % (self.Eonset,self.Etau,self.Egmax,self.Erev))
        ofile.write("%f\t%f\t%f\t%f\n" % (self.Ionset,self.Itau,self.Igmax,self.Irev))
        ofile.write("%f\t%f\n" % (self.gpas,holding_potential))
    
    def pick_randoms(self):
        self.Eonset = random.choice(self.Eonsets)
        self.Etau = random.choice(self.Etaus)
        self.Egmax = random.choice(self.Egmaxs)
        self.Erev = random.choice(self.Erevs)
        self.Ionset = random.choice(self.Ionsets)
        self.Itau = random.choice(self.Itaus)
        self.Igmax = random.choice(self.Igmaxs)
        self.Irev = random.choice(self.Irevs)
        self.gpas = random.choice(self.gpass)

    def step_one_parameter(self,parameter_number,u_or_d="u"):
        # take one parameter and step it up or down by 1/30 of the window allowed
        #   it is possible to step outside the window, but the window defines the size
        #   of the steps.
        if parameter_number == 1:
            if "u" in u_or_d:
                self.Eonset += self.EonsetW / 30.0
            else:
                self.Eonset -= self.EonsetW / 30.0
        elif parameter_number == 2:
            if "u" in u_or_d:
                self.Etau += self.EtauW / 30.0
            else:
                self.Etau -= self.EtauW / 30.0
        elif parameter_number == 3:
            if "u" in u_or_d:
                self.Egmax += self.EgmaxW / 30.0
            else:
                self.Egmax -= self.EgmaxW / 30.0
        elif parameter_number == 4:
            if "u" in u_or_d:
                self.Erev += self.ErevW / 30.0
            else:
                self.Erev -= self.ErevW / 30.0
        elif parameter_number == 5:
            if "u" in u_or_d:
                self.Ionset += self.IonsetW / 30.0
            else:
                self.Ionset -= self.IonsetW / 30.0
        elif parameter_number == 6:
            if "u" in u_or_d:
                self.Itau += self.ItauW / 30.0
            else:
                self.Itau -= self.ItauW / 30.0
        elif parameter_number == 7:
            if "u" in u_or_d:
                self.Igmax += self.IgmaxW / 30.0
            else:
                self.Igmax -= self.IgmaxW / 30.0
        elif parameter_number == 8:
            if "u" in u_or_d:
                self.Irev += self.IrevW / 30.0
            else:
                self.Irev -= self.IrevW / 30.0
        elif parameter_number == 9:
            if "u" in u_or_d:
                self.gpas += self.gpasW / 30.0
            else:
                self.gpas -= self.gpasW / 30.0
        

    def just_print(self):
        print "Eonset = " + str(self.Eonset)
        print "Etau = " + str(self.Etau)
        print "Egmax = " + str(self.Egmax)
        print "Erev = " + str(self.Erev)
        print "Ionset = " + str(self.Ionset)
        print "Itau = " + str(self.Itau)
        print "Igmax = " + str(self.Igmax)
        print "Irev = " + str(self.Irev)
        print "gpas = " + str(self.gpas)

    def exercise(self):
        # now that we have the data in a set of lists, lets read in the config file and try some stuff out.
        self.rvects = [[],[],[],[],[]]
        self.s1vects = [[],[],[],[],[]]
        self.s2vects = [[],[],[],[],[]]

        self.residual = 0
        for hp in self.holding_potentials:
            self.neuron_print(hp)
            os.system("nrniv sim_synapses.hoc >& junk.out")
            results = open("full_output.out","r")
            rlines = results.readlines()
            self.potentialVect = rlines[0].split()
            self.s1Vect = rlines[1].split()
            self.s2Vect = rlines[2].split()
            self.rvects[self.holding_potentials.index(hp)] = copy.deepcopy(self.potentialVect)
            self.s1vects[self.holding_potentials.index(hp)] = copy.deepcopy(self.s1Vect)
            self.s2vects[self.holding_potentials.index(hp)] = copy.deepcopy(self.s2Vect)
        #self.just_print()
        #self.pylab_print()

    def pylab_print(self, print_trials=False):
        pylab.figure(self.image_counter)
        self.image_counter += 1
        for i in range(len(self.data)):
            pylab.subplot(len(self.data),1,i+1)
            pylab.plot(self.time,self.data[i].average,self.time,self.rvects[i])
        if self.image_counter == 1:
            pylab.show()
        if print_trials:
            pylab.figure(99)
            for i in range(len(self.data)):
                pylab.subplot(len(self.data),1,i+1)
                pylab.plot(self.time,self.data[i].trace[0],self.time,\
                            self.data[i].trace[1],self.time,self.data[i].trace[2])

    def gnuplot_print(self):
        ofile = open("gnuplot.out", 'w')
        for ti in range(len(self.time)):
            ofile.write("%f\t" % (self.time[ti]))
            for hpi in range(len(self.holding_potentials)):
                ofile.write("%f\t%s\t%s\t%s\t" % (self.data[hpi].average[ti],self.rvects[hpi][ti],\
                                              self.s1vects[hpi][ti],self.s2vects[hpi][ti] ))
            ofile.write("\n")




class traces:
    def __init__(self,column_list,line_list):
        self.trace = []
        for i in range(len(column_list)):
            self.trace.append([])
        for line in line_list:
            tokens = line.split()
            for col in column_list:
                self.trace[column_list.index(col)].append(float(tokens[col]))
        self.calc_average_trace()
        
    def calc_average_trace(self):
        self.average = []
        for i in range(len(self.trace[0])):
            tsum = 0
            for trace in self.trace:
                tsum += trace[i]
            tsum /= len(self.trace)
            self.average.append(tsum)
        self.average = running_average(self.average,10) # 10 pt running avg


sinfo = synapse_info()
# handfit values here
sinfo.Eonset = 15.83
sinfo.Etau   = 8.1
sinfo.Egmax  = 0.118
sinfo.Erev   = -28.7
sinfo.Ionset = 13.7
sinfo.Itau   = 23.0
sinfo.Igmax  = 0.517
sinfo.Irev   = -73.33
sinfo.gpas   = 0.024

sinfo.just_print()
sinfo.exercise()
sinfo.gnuplot_print()
    
